from enum import Enum
from pathlib import Path
from gooey import Gooey, GooeyParser
from tools import create_dataset, dataset_enhancment, json2yolo

cwd = str(Path().cwd())


class Task(Enum):
    BOX_TO_YOLO = ("labelme to yolo(box)",)
    SPLIT = ("dataset split",)
    ENHANCE = ("dataset enhancement",)


@Gooey(language="chinese")
def main():
    parser = GooeyParser()
    page_parser = parser.add_subparsers(dest="JOB")

    _json2yolo(page_parser.add_parser(Task.BOX_TO_YOLO.value))
    _data_split(page_parser.add_parser(Task.SPLIT.value))
    _enhandcement(page_parser.add_parser(Task.ENHANCE.value))

    args = parser.parse_args()

    task_type = Task(args.JOB)
    if task_type == Task.BOX_TO_YOLO:
        _process_json2yolo(args)
    elif task_type == Task.SPLIT:
        _process_split(args)
    elif task_type == Task.ENHANCE:
        _process_enhance(args)


def _process_enhance(args):
    print("开始数据增强...")
    img_path, label_path = Path(args.img), Path(args.label)
    if not img_path.exists():
        raise FileNotFoundError(str(img_path))
    if not label_path.exists():
        raise FileNotFoundError(str(label_path))

    dataset_enhancment.main(img_path, label_path)
    print("增强结束")


def _process_split(args):
    print("开始拆分数据...")
    create_dataset.split_data(
        args.img,
        args.yolo,
        float(args.train_radio),
        float(args.test_radio),
        args.splits_output if args.splits_output else ".",
    )
    print("拆分结束")


def _process_json2yolo(args):
    print("开始将josn文件转为yolo格式...")
    yolo_path = Path(args.yolo)
    json2yolo.main(Path(args.json), yolo_path)
    if args.visual:
        target = Path(args.visual_dir)
        target.mkdir(exist_ok=True, parents=True)
        print("可视化结果将输出到：", target)
        json2yolo.visual_by_yolo(Path(args.img), yolo_path, target)
        print("可视化结束")
    print("转换结束")


def _json2yolo(parser):
    parser.add_argument(
        "img",
        metavar="图像文件夹",
        default=cwd,
        help="图像所在文件夹",
        widget="DirChooser",
    )
    parser.add_argument(
        "yolo",
        metavar="yolo txt文件夹",
        default=cwd,
        help="json转换后的文件所在文件夹",
        widget="DirChooser",
    )
    parser.add_argument(
        "json",
        metavar="json文件夹",
        default=cwd,
        help="json文件所在文件夹",
        widget="DirChooser",
    )
    parser.add_argument(
        "--visual",
        metavar="是否结果可视化",
        widget="CheckBox",
        help="如果选择，则会将标注后的框画到图像上面",
        required=False,
        action="store_true",
        default=False,
    )
    parser.add_argument(
        "--visual_dir",
        widget="DirChooser",
        metavar="可视化结果输出文件夹",
        help="如果选择了要结果可视化，请选择路径，默认为当前工作目录",
        default=cwd,
        required=False,
    )


def _data_split(parser):
    parser.add_argument(
        "img",
        metavar="图像文件夹",
        default=cwd,
        help="图像所在文件夹",
        widget="DirChooser",
    )
    parser.add_argument(
        "yolo",
        metavar="yolo txt文件夹",
        default=cwd,
        help="json转换后的文件所在文件夹",
        widget="DirChooser",
    )
    parser.add_argument(
        "splits_output",
        metavar="拆分输出文件夹",
        help="拆分后的输出目录",
        widget="DirChooser",
        default=cwd,
    )
    parser.add_argument(
        "train_radio",
        metavar="训练集占比",
        help="训练数据集占比",
        widget="TextField",
        default=0.7,
    )
    parser.add_argument(
        "test_radio",
        metavar="测试集占比",
        help="测试数据集占比",
        widget="TextField",
        default=0.2,
    )


def _enhandcement(parser):
    parser.add_argument(
        "img",
        metavar="图像文件根文件夹",
        help="选择包含有所有要增强图像文件的文件夹",
        widget="DirChooser",
    )
    parser.add_argument(
        "label",
        metavar="标签文件根文件夹",
        help="请注意：本文件下的文件结构应和图像文件根文件夹一致（除了文件后缀外，文件名，目录名均需要一致）",
        widget="DirChooser",
    )


main()
